# -*- coding: utf-8 -*-
"""
图像数据集加载与预处理工具包

功能：实现PyTorch数据集的索引读取、图像加载、预处理流水线
典型应用：语义分割、目标检测等计算机视觉任务
"""

import random
import cv2
import numpy as np
import torch
from torch.utils.data import Dataset


def readIndex(index_path, shuffle=False):
    """读取图像-标签路径索引文件 返回"路径对"列表

    Args:
        index_path (str): 索引文件路径，每行格式为"图像路径 标签路径"
        shuffle (bool): 是否随机打乱数据顺序，默认为False

    Returns:
        list: 路径对列表，格式为[[img_path1, label_path1], [img_path2, label_path2], ...]

    Example:
        >>> readIndex("data/train.txt", shuffle=True)
        train.txt的内容为：
        data/images/1.jpg data/labels/1.png   // 一行
        ...
        返回为：
        [['data/images/1.jpg', 'data/labels/1.png'], ...]
    """
    img_list = []
    with open(index_path, 'r') as file_to_read:
        while True:
            lines = file_to_read.readline()
            if not lines:
                break
            item = lines.strip().split()  # 分割每行为路径对
            img_list.append(item)
    file_to_read.close()

    if shuffle is True:
        random.shuffle(img_list)  # 随机打乱数据顺序
    return img_list


class dataReadPip(object):
    """
    数据读取与预处理流水线

    功能：
        1. 加载图像和标签
        2. 应用数据增强变换
        3. 执行标准化预处理

    Attributes:
        transforms (callable): 数据增强函数，需支持(img, mask)输入输出
    """

    def __init__(self, transforms=None):
        """初始化流水线

        Args:
            transforms (callable, optional): 数据增强函数，默认为None
        """
        self.transforms = transforms

    def __call__(self, item):
        """执行单样本处理流程

        Args:
            item (list): 包含图像路径和标签路径的列表

        Returns:
            tuple: 处理后的图像张量和标签张量
        """
        # 图像加载（BGR格式，0-255范围）
        img = cv2.imread(item[0])
        if img is None:  # 检查图像是否加载成功
            print(f"图像加载失败: {item[0]}")

        # 标签加载与灰度化处理
        lab = cv2.imread(item[1])
        if len(lab.shape) != 2:  # 确保标签为单通道
            lab = cv2.cvtColor(lab, cv2.COLOR_BGR2GRAY)

        # 数据增强（如翻转、旋转等）
        if self.transforms is not None:
            img, lab = self.transforms(img, lab)

        # 标准化预处理
        img = _preprocess_img(img)  # 转换为PyTorch张量
        lab = _preprocess_lab(lab)
        return img, lab


def _preprocess_img(cvImage):
    """图像标准化预处理

    Args:
        cvImage (np.ndarray): 输入图像(HWC, BGR, 0-255)

    Returns:
        torch.Tensor: 处理后的张量(CHW, BGR, float32, 0-1范围)
    """
    # 通道顺序转换 HWC -> CHW 并归一化
    cvImage = cvImage.transpose(2, 0, 1).astype(np.float32) / 255
    return torch.from_numpy(cvImage)


def _preprocess_lab(cvImage):
    """标签标准化预处理

    Args:
        cvImage (np.ndarray): 输入标签(单通道, 0或255)

    Returns:
        torch.Tensor: 二值化标签张量(0或1, float32)
    """
    cvImage = cvImage.astype(np.float32) / 255  # 归一化到[0,1]
    return torch.from_numpy(cvImage)


class loadedDataset(Dataset):
    """PyTorch标准数据集封装

    功能：
        1. 实现__len__和__getitem__接口
        2. 支持自定义预处理流水线

    Attributes:
        dataset (list): 原始数据集列表
        preprocess (callable): 预处理函数
    """

    def __init__(self, dataset, preprocess=None):
        """初始化数据集

        Args:
            dataset (list): 原始数据列表
            preprocess (callable, optional): 预处理函数，默认为恒等变换
        """
        super(loadedDataset, self).__init__()
        self.dataset = dataset
        self.preprocess = preprocess if preprocess else lambda x: x

    def __getitem__(self, index):
        """获取单样本

        Args:
            index (int): 样本索引

        Returns:
            any: 预处理后的样本数据
        """
        return self.preprocess(self.dataset[index])

    def __len__(self):
        """获取数据集长度

        Returns:
            int: 数据集样本数量
        """
        return len(self.dataset)